package me.huanmeng.heater.plugin;

import com.google.common.io.ByteStreams;
import lombok.Getter;
import me.huanmeng.heater.exception.PluginLoadException;
import org.jetbrains.annotations.Nullable;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.net.URLClassLoader;
import java.security.CodeSigner;
import java.security.CodeSource;
import java.util.Enumeration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.Manifest;

/**
 * 2021/9/29<br>
 * Heater<br>
 *
 * @author huanmeng_qwq
 */
public class PluginClassLoader extends URLClassLoader {
    private JarFile jar;
    private Manifest manifest;
    private URL url;
    @Getter
    private Map<String, Class<?>> classes = new ConcurrentHashMap<>();
    public final JavaPlugin plugin;

    public PluginClassLoader(File file,File dataFolder, PluginDescription description, ClassLoader parent) throws Exception {
        super(new URL[]{file.toURI().toURL()}, parent);
        this.jar = new JarFile(file);
        this.manifest = jar.getManifest();
        this.url = file.toURI().toURL();
        String main = description.getMain();
        try {
            Class<?> mainClass;
            try {
                mainClass = Class.forName(main, true, this);
            } catch (ClassNotFoundException ex) {
                throw new PluginLoadException("Cannot find main class `" + main + "'", ex);
            }

            Class<? extends JavaPlugin> pluginClass;
            try {
                pluginClass = mainClass.asSubclass(JavaPlugin.class);
            } catch (ClassCastException ex) {
                throw new PluginLoadException("main class `" + main + "' does not extend JavaPlugin", ex);
            }

            plugin = pluginClass.newInstance();
            plugin.initialize(file,dataFolder,description,this);
        } catch (IllegalAccessException | InstantiationException e) {
            throw new PluginLoadException(e);
        }
    }

    @Nullable
    @Override
    public URL getResource(String name) {
        return findResource(name);
    }

    @Override
    public Enumeration<URL> getResources(String name) throws IOException {
        return findResources(name);
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        Class<?> clazz = classes.get(name);
        if (clazz == null) {
            String path = name.replace('.', '/').concat(".class");
            JarEntry entry = jar.getJarEntry(path);
            if (entry != null) {
                byte[] classBytes;

                try (InputStream is = jar.getInputStream(entry)) {
                    classBytes = ByteStreams.toByteArray(is);
                } catch (IOException ex) {
                    throw new ClassNotFoundException(name, ex);
                }

                int dot = name.lastIndexOf('.');
                if (dot != -1) {
                    String pkgName = name.substring(0, dot);
                    if (getPackage(pkgName) == null) {
                        try {
                            if (manifest != null) {
                                definePackage(pkgName, manifest, url);
                            } else {
                                definePackage(pkgName, null, null, null, null, null, null, null);
                            }
                        } catch (IllegalArgumentException ex) {
                            if (getPackage(pkgName) == null) {
                                throw new IllegalStateException("Cannot find package " + pkgName);
                            }
                        }
                    }
                }

                CodeSigner[] signers = entry.getCodeSigners();
                CodeSource source = new CodeSource(url, signers);

                clazz = defineClass(name, classBytes, 0, classBytes.length, source);
            } else {
                clazz = super.findClass(name);
            }
        }
        classes.put(name, clazz);
        return clazz;
    }

    @Override
    public void close() throws IOException {
        try {
            super.close();
        } finally {
            jar.close();
        }
    }

    public JavaPlugin getPlugin() {
        return plugin;
    }
}
